#! /usr/bin/env python3
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import pickle
from PIL import Image, ImageDraw
import numpy as np
import os

def visualization(data_path):
    with open(data_path, 'rb') as f:
        fbs = pickle.load(f)
    for row in fbs:
        filename = row[0]
        boxs = row[1]
        img = Image.open(filename)

        fig, ax = plt.subplots(figsize=(12, 6))
        ax.set_title(os.path.basename(filename))
        ax.imshow(np.asarray(img, dtype='uint8'))
        for key in boxs.keys():
            if key == 'L1':
                edgecolor = 'orange'
            else:
                edgecolor = 'red'
            for box in boxs[key]:
                [[minr, minc], [maxr, maxc]] = box
                rect = mpatches.Rectangle((minc, minr), maxc - minc, maxr - minr,
                                          fill=False, edgecolor=edgecolor, linewidth=2)
                ax.add_patch(rect)
        plt.show()
    print('finished!')

if __name__ == '__main__': visualization('./fbs.pkl')
